# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: yuxuan

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

import os
import datetime
import glob
from collections import defaultdict

import numpy as np
import pandas as pd
from absl import (flags, app)

FLAGS = flags.FLAGS

flags.DEFINE_string('include', '', '')
flags.DEFINE_string('exclude', '', '')
flags.DEFINE_string('cutoff_hour_min', '', '')
flags.DEFINE_float('duration_hours', 24, '')
flags.DEFINE_string('strategy_name', '', '')
flags.DEFINE_string('output_dir', 'plot/spread_basis', '')


def find_all_balance_within(from_dt, to_dt, log_paths, includes, excludes):
  ratios = defaultdict(list)
  for log_path in log_paths:
    with open(log_path, 'r') as log_file:
      in_log = False
      for line in log_file:
        line = line.strip()
        if line.startswith('*'):
          in_log = True
          fmt = '%Y-%m-%d %H:%M:%S.%f' if '.' in line else '%Y-%m-%d %H:%M:%S'
          dt = datetime.datetime.strptime(line.split(' ', 1)[1], fmt)
          if dt < from_dt:
            in_log = False
            continue
          elif dt > to_dt:
            break
        elif in_log:
          if ':' in line:
            in_log = False
          elif 'Currency' in line or '---' in line:
            continue
          else:
            data = [item for item in line.split(' ') if item]
            ticker = data[0].split('.')[0]
            pos0 = float(data[1])
            pos1 = float(data[2])
            reserve = float(data[4])
            if reserve == 0:
              continue
            if (pos0 + pos1) == 0:
              ratio = np.nan
            else:
              ratio = pos0 / (pos0 + pos1)
            ratios[ticker].append([dt, ratio, (pos0 + pos1) / reserve - 1])
  df = None
  for ticker, a_ratios in ratios.items():
    if includes and ticker not in includes:
      continue
    if excludes and ticker in excludes:
      continue
    a_df = pd.DataFrame(a_ratios, columns=['dt', ticker, '%s_mb' % ticker])
    a_df = a_df.set_index('dt')
    df = a_df if df is None else \
        df.merge(a_df, left_index=True, right_index=True)
  tk_columns = list(ratios.keys())
  mb_columns = [c for c in df.columns if c not in tk_columns]
  return df[tk_columns], df[mb_columns]


def find_start_end_time(cutoff_hour_min, duration_hours):
  now = datetime.datetime.now()
  dt = now.strftime('%Y%m%d')
  to_dt = datetime.datetime.strptime('%s %s' % (dt, cutoff_hour_min), '%Y%m%d %H:%M')
  if to_dt > now:
    to_dt -= datetime.timedelta(days=1)
  from_dt = to_dt - datetime.timedelta(hours=float(duration_hours))
  return from_dt, to_dt


def find_all_logs(strategy_name, start_time, end_time, archive_dir='archive'):
  archive_logs = sorted(glob.glob('%s/log_%s*/*' % (archive_dir, strategy_name)))
  start_dt = int(start_time.strftime('%Y%m%d'))
  archive_logs = [
      log for log in archive_logs
      if int(os.path.dirname(log).split('_')[-1].split('-')[0]) >= start_dt
  ]
  current_log = glob.glob('log_%s/*' % strategy_name)
  return archive_logs + current_log


def main(argv):
  if not FLAGS.cutoff_hour_min:
    return 1
  if not FLAGS.strategy_name:
    return 1
  includes = set()
  if FLAGS.include:
    includes = set([ticker.strip() for ticker in FLAGS.include.split(',')])
  excludes = set()
  if FLAGS.exclude:
    excludes = set([ticker.strip() for ticker in FLAGS.exclude.split(',')])

  from_dt, to_dt = find_start_end_time(FLAGS.cutoff_hour_min, FLAGS.duration_hours)
  logs = find_all_logs(FLAGS.strategy_name, from_dt, to_dt)
  tk_df, mb_df = find_all_balance_within(from_dt, to_dt, logs, includes, excludes)
  fig, (ax, ax2) = plt.subplots(nrows=2, sharex=True, figsize=(12, 12))

  tk_df.plot(lw=1, ylim=(0, 1), ax=ax)
  mb_df.plot(lw=1, ax=ax2)
  from_dt = from_dt.strftime('%Y%m%d%H%M%S')
  to_dt = to_dt.strftime('%Y%m%d%H%M%S')
  if not os.path.exists(FLAGS.output_dir):
    os.makedirs(FLAGS.output_dir)
  plt.savefig(os.path.join(FLAGS.output_dir, '%s_%s_%s.png' %
                           (FLAGS.strategy_name, from_dt, to_dt)))


if __name__ == '__main__':
  app.run(main)
